#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/12/16 14:06
# @Author  : yingxiao zhang
# @File    : stock_ana.py

import pandas as pd
import glob
import json

def calculate_max_drawdown(prices):
    # 初始化最大回撤和峰值
    max_drawdown = 0
    peak = prices[0]

    for price in prices:
        # 更新峰值
        if price > peak:
            peak = price
        # 计算当前回撤
        drawdown = (peak - price) / peak
        # 更新最大回撤
        if drawdown > max_drawdown:
            max_drawdown = drawdown

    return max_drawdown

files = glob.glob('stocks_data/*.csv')
stocks_data = pd.read_csv('data/stock_list.csv')

# choosed_stocks = dict()
choosed_stocks = []
for file in files:
    data = pd.read_csv(file)
    ts_code = data['ts_code'][0]
    prices = data['close'].tolist()
    prices.reverse()
    max_drawdown = calculate_max_drawdown(prices)
    ayear = 242
    if max_drawdown < 0.6 or ts_code in {'002595.SZ', '601838.SH'}:
        tenyear_grow, fiveyear_grow, threeyear_grow, oneyear_grow = 0, 0, 0, 0
        if (len(data) >= 250 * 10):
            tenyear_grow = (data['close'][0] - data['close'][ayear*10 - 1]) / data['close'][ayear*10 - 1]
        if (len(data) >= 250 * 5):
            fiveyear_grow = (data['close'][0] - data['close'][ayear * 5 - 1]) / data['close'][ayear * 5 - 1]
        if (len(data) >= 250 * 3):
            threeyear_grow = (data['close'][0] - data['close'][ayear * 3 - 1]) / data['close'][ayear * 3 - 1]
        if (len(data) >= 250 * 1):
            oneyear_grow = (data['close'][0] - data['close'][ayear * 1 - 1]) / data['close'][ayear * 1 - 1]

        if ts_code in {'002595.SZ', '601838.SH'}:
            stock_name = stocks_data.loc[stocks_data['ts_code'] == ts_code, 'name'].iloc[0]
            choosed_stocks.append(
                {'ts_code': ts_code, 'name': stock_name, '最大回撤': max_drawdown, '一年收益率': oneyear_grow,
                 '三年收益率': threeyear_grow, '五年收益率': fiveyear_grow, '十年收益率': tenyear_grow})
            continue

        if oneyear_grow < 0.05 or threeyear_grow < 0.15 or fiveyear_grow < 0.25 :
            continue
        if (threeyear_grow-oneyear_grow) < 0.1 or (fiveyear_grow-threeyear_grow) < 0.1:
            continue
        stock_name = stocks_data.loc[stocks_data['ts_code'] == ts_code, 'name'].iloc[0]
#         stock_data = {'name': stock_name, '最大回撤': max_drawdown, '一年收益率': oneyear_grow, '三年收益率': threeyear_grow, '五年收益率': fiveyear_grow, '十年收益率': tenyear_grow}
#         print(ts_code, stock_data)
#         choosed_stocks[ts_code] = stock_data

        choosed_stocks.append({'ts_code': ts_code,'name': stock_name, '最大回撤': max_drawdown, '一年收益率': oneyear_grow, '三年收益率': threeyear_grow, '五年收益率': fiveyear_grow, '十年收益率': tenyear_grow})
#
# # 将字典写入JSON文件
# with open('choosed_stocks.json', 'w', encoding='utf-8') as json_file:
#     json.dump(choosed_stocks, json_file, ensure_ascii=False, indent=4)

df = pd.DataFrame(choosed_stocks)

# 将 DataFrame 写入 Excel 文件
df.to_excel('choosed_stocks.xlsx', index=False)

print(len(choosed_stocks))

